#!/usr/bin/env python

import sys
import array

from util import *
from bwt_enc import bwt, ibwt
from runlen_enc import runlen_enc, runlen_dec
from huffman_enc import *

chunklen = 512
cutoff = 5000000
reverselist = False

class FinalEncoder:

    def __init__(self):

        self.sym_cnt = {}
        self.len_cnt = {}
        self.num_runs = 0

        self.sym_code = {}
        self.len_code = {}

    def processchunk(self, chunk, phase):
        if reverselist:
            chunk.reverse()
        
        # bwt encoding first
        chunk_bwt_coded = bwt(chunk)

        # runlen encoding then
        chunk_runlen_coded = runlen_enc(chunk_bwt_coded)

        self.num_runs += len(chunk_runlen_coded)

        if phase == 1:
            for (i, s) in chunk_runlen_coded:
                try:
                    self.sym_cnt[s] += 1
                except:
                    self.sym_cnt[s] = 1

                try:
                    self.len_cnt[i] += 1
                except:
                    self.len_cnt[i] = 1

        elif phase == 2:

            for (i, s) in chunk_runlen_coded:
                self.total_sym += i
                self.enc_buf += self.len_code[str(i)] + self.sym_code[str(s)]

                while (len(self.enc_buf) >= 8):
                    self.binary.append(int(self.enc_buf[:8], 2))
                    self.enc_buf = self.enc_buf[8:]

            

    def enc(self, rawlist):

        self.total_sym = 0
        self.binary = array.array('B')
        self.enc_buf = ""

        chunk = []
        step = 1
        s0 = None
        for l in rawlist:
            if s0 == None:
                s0 = l
                continue
            chunk.append(l)
            if len(chunk) == chunklen:
                #print "phase1", step
                if step == cutoff:
                    break
                self.processchunk(chunk, 1)
                chunk = []
                step += 1

        if chunk:
            self.processchunk(chunk, 1)

        runlen_sym_dist = []
        runlen_len_dist = []

        for (k,v) in self.sym_cnt.items():
            runlen_sym_dist.append((1.0 * v / self.num_runs, str(k)))

        for (k,v) in self.len_cnt.items():
            runlen_len_dist.append((1.0 * v / self.num_runs, str(k)))

        getHuffCode(makeHuffTree(runlen_sym_dist), self.sym_code)
        getHuffCode(makeHuffTree(runlen_len_dist), self.len_code)


        chunk = []
        step = 1
        s0 = None
        for l in rawlist:
            if s0 == None:
                s0 = l
                continue
            chunk.append(l)
            if len(chunk) == chunklen:
                #print "phase2", step
                if step == cutoff:
                    break
                coded = self.processchunk(chunk, 2)
                chunk = []
                step += 1
        if chunk:
            self.processchunk(chunk, 2)

        if (len(self.enc_buf) > 0):
            assert (len(self.enc_buf) < 8)
            self.binary.append(int(self.enc_buf+"0"*(8 - len(self.enc_buf)), 2) )

        return [s0, inv_dict(self.sym_code), inv_dict(self.len_code), self.total_sym, self.binary]



    def next_bit(self, binary):

        if self.dec_buf == "":
            self.dec_buf += "{0:08b}".format(binary[self.dec_idx])
            self.dec_idx += 1
        ret = self.dec_buf[0]
        self.dec_buf = self.dec_buf[1:]
        return ret

    def dec(self, coded):

        [s0, inv_sym_code, inv_len_code, total_sym, binary] = coded

        print "total_sym = ", total_sym
        final_code = [s0]
        num_sym = 0
        self.dec_buf = ""
        self.dec_idx = 0

        while num_sym < total_sym:
            nc = 0
            runlen_code = []
            while nc < (chunklen+1) and num_sym < total_sym:
                next_code = ""
                while not(next_code in inv_len_code.keys()):
                    next_code += self.next_bit(binary)
                r = int(inv_len_code[next_code])
                next_code = ""
                while not(next_code in inv_sym_code.keys()):
                    next_code += self.next_bit(binary)
                s = inv_sym_code[next_code]
                runlen_code.append((r,s))
                nc += r
                num_sym += r

            c1 = runlen_dec(runlen_code)
            
            c2 = ibwt(c1)
            
            if reverselist:
                c2.reverse()
            final_code += c2
        return final_code
            
    


if __name__ == "__main__":
    input = ["10", "2", "11", "6", "11",  "6"]
    Enc = FinalEncoder()
    
    c = Enc.enc(input)
    
    print "input=", input

    print "after final_dec", Enc.dec(c)

    # for metric in ['temp2C', 'presure', 'humidity']:
    #     print metric
    #     input = []
    #     for l in open ("delta_seglen-all_%s.txt"%(metric)):
    #         if l.startswith("#"): continue
    #         input.append(eval(l))
    #     final_enc(input)
